#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
API服务集成模块

展示如何将新的LangChain代理系统集成到现有的API服务中
"""

from fastapi import FastAPI, HTTPException, BackgroundTasks
from pydantic import BaseModel
from typing import Dict, Any, Optional
import uuid
from datetime import datetime, timedelta
from loguru import logger
import asyncio
from contextlib import asynccontextmanager

# 导入代理模块
from __init__ import FloodForecastAgent, create_agent

# 会话管理
class SessionManager:
    """会话管理器"""
    
    def __init__(self, session_timeout: int = 3600):
        self.sessions: Dict[str, Dict[str, Any]] = {}
        self.session_timeout = session_timeout  # 会话超时时间（秒）
    
    def create_session(self, llm_config: Dict[str, Any]) -> str:
        """创建新会话"""
        session_id = str(uuid.uuid4())
        agent = create_agent(llm_config)
        
        self.sessions[session_id] = {
            "agent": agent,
            "created_at": datetime.now(),
            "last_activity": datetime.now(),
            "llm_config": llm_config
        }
        
        logger.info(f"创建新会话: {session_id}")
        return session_id
    
    def get_session(self, session_id: str) -> Optional[FloodForecastAgent]:
        """获取会话"""
        if session_id not in self.sessions:
            return None
        
        session = self.sessions[session_id]
        
        # 检查会话是否过期
        if datetime.now() - session["last_activity"] > timedelta(seconds=self.session_timeout):
            self.delete_session(session_id)
            return None
        
        # 更新最后活动时间
        session["last_activity"] = datetime.now()
        return session["agent"]
    
    def delete_session(self, session_id: str):
        """删除会话"""
        if session_id in self.sessions:
            del self.sessions[session_id]
            logger.info(f"删除会话: {session_id}")
    
    def cleanup_expired_sessions(self):
        """清理过期会话"""
        now = datetime.now()
        expired_sessions = []
        
        for session_id, session in self.sessions.items():
            if now - session["last_activity"] > timedelta(seconds=self.session_timeout):
                expired_sessions.append(session_id)
        
        for session_id in expired_sessions:
            self.delete_session(session_id)
        
        if expired_sessions:
            logger.info(f"清理了 {len(expired_sessions)} 个过期会话")

# 全局会话管理器
session_manager = SessionManager()

# 定期清理任务
async def cleanup_task():
    """定期清理过期会话"""
    while True:
        try:
            session_manager.cleanup_expired_sessions()
            await asyncio.sleep(300)  # 每5分钟清理一次
        except Exception as e:
            logger.error(f"清理任务出错: {str(e)}")
            await asyncio.sleep(60)  # 出错后等待1分钟再试

# FastAPI应用生命周期管理
@asynccontextmanager
async def lifespan(app: FastAPI):
    # 启动时
    logger.info("启动API服务")
    cleanup_task_handle = asyncio.create_task(cleanup_task())
    
    yield
    
    # 关闭时
    logger.info("关闭API服务")
    cleanup_task_handle.cancel()
    try:
        await cleanup_task_handle
    except asyncio.CancelledError:
        pass

# 创建FastAPI应用
app = FastAPI(
    title="洪水预报智能代理API",
    description="基于LangChain的洪水预报参数收集和验证服务",
    version="1.0.0",
    lifespan=lifespan
)

# 请求模型
class CreateSessionRequest(BaseModel):
    """创建会话请求"""
    llm_config: Dict[str, Any]

class ChatRequest(BaseModel):
    """对话请求"""
    session_id: str
    message: str

class SessionSummaryRequest(BaseModel):
    """会话摘要请求"""
    session_id: str

# API端点
@app.post("/api/v1/sessions")
async def create_session(request: CreateSessionRequest):
    """创建新的对话会话"""
    try:
        session_id = session_manager.create_session(request.llm_config)
        return {
            "success": True,
            "session_id": session_id,
            "message": "会话创建成功",
            "timestamp": datetime.now().isoformat()
        }
    except Exception as e:
        logger.error(f"创建会话失败: {str(e)}")
        raise HTTPException(status_code=500, detail=f"创建会话失败: {str(e)}")

@app.post("/api/v1/chat")
async def chat(request: ChatRequest):
    """处理对话消息"""
    try:
        agent = session_manager.get_session(request.session_id)
        if not agent:
            raise HTTPException(status_code=404, detail="会话不存在或已过期")
        
        # 处理用户输入
        response = agent.process_user_input(request.message)
        
        return {
            "success": True,
            "data": response,
            "timestamp": datetime.now().isoformat()
        }
        
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"处理对话失败: {str(e)}")
        raise HTTPException(status_code=500, detail=f"处理对话失败: {str(e)}")

@app.get("/api/v1/sessions/{session_id}/summary")
async def get_session_summary(session_id: str):
    """获取会话摘要"""
    try:
        agent = session_manager.get_session(session_id)
        if not agent:
            raise HTTPException(status_code=404, detail="会话不存在或已过期")
        
        summary = agent.get_session_summary()
        
        return {
            "success": True,
            "data": summary,
            "timestamp": datetime.now().isoformat()
        }
        
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"获取会话摘要失败: {str(e)}")
        raise HTTPException(status_code=500, detail=f"获取会话摘要失败: {str(e)}")

@app.delete("/api/v1/sessions/{session_id}")
async def delete_session(session_id: str):
    """删除会话"""
    try:
        agent = session_manager.get_session(session_id)
        if not agent:
            raise HTTPException(status_code=404, detail="会话不存在")
        
        session_manager.delete_session(session_id)
        
        return {
            "success": True,
            "message": "会话删除成功",
            "timestamp": datetime.now().isoformat()
        }
        
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"删除会话失败: {str(e)}")
        raise HTTPException(status_code=500, detail=f"删除会话失败: {str(e)}")

@app.post("/api/v1/sessions/{session_id}/reset")
async def reset_session(session_id: str):
    """重置会话状态"""
    try:
        agent = session_manager.get_session(session_id)
        if not agent:
            raise HTTPException(status_code=404, detail="会话不存在或已过期")
        
        agent.reset_session()
        
        return {
            "success": True,
            "message": "会话重置成功",
            "timestamp": datetime.now().isoformat()
        }
        
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"重置会话失败: {str(e)}")
        raise HTTPException(status_code=500, detail=f"重置会话失败: {str(e)}")

@app.get("/api/v1/health")
async def health_check():
    """健康检查"""
    return {
        "status": "healthy",
        "timestamp": datetime.now().isoformat(),
        "active_sessions": len(session_manager.sessions)
    }

@app.get("/api/v1/stats")
async def get_stats():
    """获取服务统计信息"""
    return {
        "active_sessions": len(session_manager.sessions),
        "session_timeout": session_manager.session_timeout,
        "timestamp": datetime.now().isoformat()
    }

# 兼容性端点 - 与原有API保持兼容
@app.post("/collect_parameters")
async def collect_parameters_legacy(request: Dict[str, Any]):
    """兼容原有API的参数收集端点"""
    try:
        # 从请求中提取LLM配置
        llm_config = request.get("llm_config", {
            "model": "qwq:latest",
            "temperature": 0.1
        })
        
        # 创建临时会话
        session_id = session_manager.create_session(llm_config)
        agent = session_manager.get_session(session_id)
        
        # 处理用户输入
        user_input = request.get("user_input", "")
        response = agent.process_user_input(user_input)
        
        # 清理临时会话（如果已完成）
        if response.get("is_complete", False):
            session_manager.delete_session(session_id)
        
        return {
            "success": True,
            "response": response["response"],
            "collected_parameters": response["collected_parameters"],
            "is_complete": response["is_complete"],
            "session_id": session_id if not response.get("is_complete", False) else None
        }
        
    except Exception as e:
        logger.error(f"兼容性API调用失败: {str(e)}")
        raise HTTPException(status_code=500, detail=f"处理请求失败: {str(e)}")

if __name__ == "__main__":
    import uvicorn
    
    # 配置日志
    logger.add("logs/agent_api.log", rotation="1 day", retention="30 days")
    
    # 启动服务
    uvicorn.run(
        "api_integration:app",
        host="0.0.0.0",
        port=8000,
        reload=True,
        log_level="info"
    )